"""
工具集成器

将新的分析工具集成到现有Agent Laboratory工具链中，包括：
- 工具注册和发现
- 工具间数据传递
- 工具兼容性验证
- 工具使用统计和优化
"""
import json
import asyncio
from typing import Dict, List, Any, Optional, Callable, Union
from datetime import datetime, timedelta
from dataclasses import dataclass, field

# 导入现有工具
from .common import HFDataSearch, SemanticScholarSearch, ArxivSearch, execute_code
from .enhanced_research_tools import (
    EnhancedLiteratureAnalyzer, SmartExperimentDesigner, 
    ResearchContext
)


@dataclass
class ToolMetadata:
    """工具元数据"""
    name: str
    version: str
    description: str
    input_types: List[str]
    output_types: List[str]
    dependencies: List[str] = field(default_factory=list)
    capabilities: List[str] = field(default_factory=list)
    performance_metrics: Dict[str, float] = field(default_factory=dict)
    usage_count: int = 0
    last_used: Optional[datetime] = None


@dataclass
class ToolChain:
    """工具链"""
    chain_id: str
    name: str
    tools: List[str] = field(default_factory=list)
    data_flow: Dict[str, str] = field(default_factory=dict)  # tool_output -> tool_input
    success_rate: float = 0.0
    average_execution_time: float = 0.0


class ToolIntegrator:
    """工具集成器"""
    
    def __init__(self):
        # 工具注册表
        self.registered_tools: Dict[str, Any] = {}
        self.tool_metadata: Dict[str, ToolMetadata] = {}
        
        # 工具链管理
        self.tool_chains: Dict[str, ToolChain] = {}
        self.execution_history: List[Dict[str, Any]] = []
        
        # 数据转换器
        self.data_transformers: Dict[str, Callable] = {}
        
        # 性能监控
        self.performance_stats = {
            'total_executions': 0,
            'successful_executions': 0,
            'average_execution_time': 0.0,
            'tool_usage_frequency': {},
            'error_rates': {}
        }
        
        # 初始化核心工具
        self._register_core_tools()
        self._register_enhanced_tools()
        self._setup_data_transformers()
        self._create_default_tool_chains()
    
    def _register_core_tools(self):
        """注册核心工具"""
        # 注册现有工具
        self.register_tool(
            name="hf_data_search",
            tool_instance=HFDataSearch(),
            metadata=ToolMetadata(
                name="HuggingFace Dataset Search",
                version="1.0",
                description="搜索和推荐HuggingFace数据集",
                input_types=["query_string"],
                output_types=["dataset_list"],
                capabilities=["dataset_search", "dataset_filtering", "metadata_extraction"]
            )
        )
        
        self.register_tool(
            name="semantic_scholar_search",
            tool_instance=SemanticScholarSearch(),
            metadata=ToolMetadata(
                name="Semantic Scholar Search",
                version="1.0",
                description="搜索学术论文和获取引用信息",
                input_types=["query_string"],
                output_types=["paper_list"],
                capabilities=["paper_search", "citation_analysis", "metadata_extraction"]
            )
        )
        
        self.register_tool(
            name="arxiv_search",
            tool_instance=ArxivSearch(),
            metadata=ToolMetadata(
                name="ArXiv Search",
                version="1.0",
                description="搜索ArXiv论文和获取全文",
                input_types=["query_string", "paper_id"],
                output_types=["paper_list", "full_text"],
                capabilities=["paper_search", "full_text_extraction", "metadata_extraction"]
            )
        )
        
        self.register_tool(
            name="code_executor",
            tool_instance=execute_code,
            metadata=ToolMetadata(
                name="Code Executor",
                version="1.0",
                description="执行Python代码并返回结果",
                input_types=["code_string"],
                output_types=["execution_result"],
                capabilities=["code_execution", "result_capture", "error_handling"]
            )
        )
    
    def _register_enhanced_tools(self):
        """注册增强工具"""
        self.register_tool(
            name="enhanced_literature_analyzer",
            tool_instance=EnhancedLiteratureAnalyzer(),
            metadata=ToolMetadata(
                name="Enhanced Literature Analyzer",
                version="1.0",
                description="基于上下文的智能文献分析",
                input_types=["query_string", "research_context"],
                output_types=["literature_analysis"],
                capabilities=["context_aware_search", "gap_analysis", "trend_extraction", "citation_network"]
            )
        )
        
        self.register_tool(
            name="smart_experiment_designer",
            tool_instance=SmartExperimentDesigner(),
            metadata=ToolMetadata(
                name="Smart Experiment Designer",
                version="1.0",
                description="智能实验设计和优化",
                input_types=["research_context", "literature_analysis"],
                output_types=["experiment_design"],
                capabilities=["experiment_design", "parameter_optimization", "resource_estimation", "risk_assessment"]
            )
        )
    
    def register_tool(self, name: str, tool_instance: Any, metadata: ToolMetadata):
        """注册工具"""
        self.registered_tools[name] = tool_instance
        self.tool_metadata[name] = metadata
        self.performance_stats['tool_usage_frequency'][name] = 0
        self.performance_stats['error_rates'][name] = 0.0
        
        print(f"已注册工具: {name} - {metadata.description}")
    
    def _setup_data_transformers(self):
        """设置数据转换器"""
        # 字符串到研究上下文的转换器
        self.data_transformers['string_to_research_context'] = self._string_to_research_context
        
        # 论文列表到文献分析的转换器
        self.data_transformers['paper_list_to_literature_analysis'] = self._paper_list_to_literature_analysis
        
        # 数据集列表到数据需求的转换器
        self.data_transformers['dataset_list_to_data_requirements'] = self._dataset_list_to_data_requirements
    
    def _create_default_tool_chains(self):
        """创建默认工具链"""
        # 文献综述工具链
        self.create_tool_chain(
            chain_id="literature_review_chain",
            name="文献综述工具链",
            tools=["enhanced_literature_analyzer", "arxiv_search", "semantic_scholar_search"],
            data_flow={
                "enhanced_literature_analyzer": "arxiv_search",
                "arxiv_search": "semantic_scholar_search"
            }
        )
        
        # 实验设计工具链
        self.create_tool_chain(
            chain_id="experiment_design_chain",
            name="实验设计工具链",
            tools=["enhanced_literature_analyzer", "smart_experiment_designer", "hf_data_search"],
            data_flow={
                "enhanced_literature_analyzer": "smart_experiment_designer",
                "smart_experiment_designer": "hf_data_search"
            }
        )
        
        # 数据分析工具链
        self.create_tool_chain(
            chain_id="data_analysis_chain",
            name="数据分析工具链",
            tools=["hf_data_search", "code_executor"],
            data_flow={
                "hf_data_search": "code_executor"
            }
        )
    
    def create_tool_chain(self, chain_id: str, name: str, tools: List[str], 
                         data_flow: Dict[str, str]) -> bool:
        """创建工具链"""
        try:
            # 验证工具存在
            for tool in tools:
                if tool not in self.registered_tools:
                    raise ValueError(f"工具 {tool} 未注册")
            
            # 验证数据流
            for source, target in data_flow.items():
                if source not in tools or target not in tools:
                    raise ValueError(f"数据流中的工具 {source} -> {target} 不在工具列表中")
            
            tool_chain = ToolChain(
                chain_id=chain_id,
                name=name,
                tools=tools,
                data_flow=data_flow
            )
            
            self.tool_chains[chain_id] = tool_chain
            print(f"已创建工具链: {name}")
            return True
            
        except Exception as e:
            print(f"创建工具链失败: {e}")
            return False
    
    async def execute_tool(self, tool_name: str, input_data: Any, 
                          context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """执行单个工具"""
        start_time = datetime.now()
        
        try:
            if tool_name not in self.registered_tools:
                raise ValueError(f"工具 {tool_name} 未注册")
            
            tool_instance = self.registered_tools[tool_name]
            metadata = self.tool_metadata[tool_name]
            
            # 更新使用统计
            metadata.usage_count += 1
            metadata.last_used = start_time
            self.performance_stats['tool_usage_frequency'][tool_name] += 1
            
            # 执行工具
            result = await self._execute_tool_method(tool_instance, tool_name, input_data, context)
            
            # 计算执行时间
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # 更新性能指标
            metadata.performance_metrics['last_execution_time'] = execution_time
            if 'average_execution_time' not in metadata.performance_metrics:
                metadata.performance_metrics['average_execution_time'] = execution_time
            else:
                metadata.performance_metrics['average_execution_time'] = (
                    metadata.performance_metrics['average_execution_time'] * 0.9 + 
                    execution_time * 0.1
                )
            
            # 记录执行历史
            execution_record = {
                'tool_name': tool_name,
                'timestamp': start_time.isoformat(),
                'execution_time': execution_time,
                'success': True,
                'input_size': len(str(input_data)),
                'output_size': len(str(result))
            }
            self.execution_history.append(execution_record)
            
            # 更新全局统计
            self.performance_stats['total_executions'] += 1
            self.performance_stats['successful_executions'] += 1
            
            return {
                'success': True,
                'result': result,
                'execution_time': execution_time,
                'tool_metadata': metadata
            }
            
        except Exception as e:
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # 更新错误率
            self.performance_stats['error_rates'][tool_name] = (
                self.performance_stats['error_rates'].get(tool_name, 0.0) * 0.9 + 0.1
            )
            
            # 记录失败执行
            execution_record = {
                'tool_name': tool_name,
                'timestamp': start_time.isoformat(),
                'execution_time': execution_time,
                'success': False,
                'error': str(e)
            }
            self.execution_history.append(execution_record)
            
            self.performance_stats['total_executions'] += 1
            
            return {
                'success': False,
                'error': str(e),
                'execution_time': execution_time
            }
    
    async def _execute_tool_method(self, tool_instance: Any, tool_name: str, 
                                  input_data: Any, context: Optional[Dict[str, Any]]) -> Any:
        """执行工具方法"""
        if tool_name == "hf_data_search":
            if isinstance(input_data, str):
                return tool_instance.retrieve_ds(input_data)
            else:
                raise ValueError("HF Data Search需要字符串查询")
        
        elif tool_name == "semantic_scholar_search":
            if isinstance(input_data, str):
                return tool_instance.find_papers_by_str(input_data)
            else:
                raise ValueError("Semantic Scholar Search需要字符串查询")
        
        elif tool_name == "arxiv_search":
            if isinstance(input_data, str):
                return tool_instance.find_papers_by_str(input_data)
            elif isinstance(input_data, dict) and 'paper_id' in input_data:
                return tool_instance.retrieve_full_paper_text(input_data['paper_id'])
            else:
                raise ValueError("ArXiv Search需要字符串查询或论文ID")
        
        elif tool_name == "code_executor":
            if isinstance(input_data, str):
                return execute_code(input_data)
            else:
                raise ValueError("Code Executor需要代码字符串")
        
        elif tool_name == "enhanced_literature_analyzer":
            if isinstance(input_data, dict) and 'query' in input_data and 'context' in input_data:
                return tool_instance.analyze_literature_with_context(
                    input_data['query'], input_data['context']
                )
            else:
                raise ValueError("Enhanced Literature Analyzer需要查询和上下文")
        
        elif tool_name == "smart_experiment_designer":
            if isinstance(input_data, dict) and 'context' in input_data and 'literature_analysis' in input_data:
                return tool_instance.design_experiment(
                    input_data['context'], input_data['literature_analysis']
                )
            else:
                raise ValueError("Smart Experiment Designer需要研究上下文和文献分析")
        
        else:
            raise ValueError(f"未知工具: {tool_name}")
    
    async def execute_tool_chain(self, chain_id: str, initial_input: Any, 
                               context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """执行工具链"""
        start_time = datetime.now()
        
        try:
            if chain_id not in self.tool_chains:
                raise ValueError(f"工具链 {chain_id} 不存在")
            
            tool_chain = self.tool_chains[chain_id]
            results = {}
            current_data = initial_input
            
            # 按工具链顺序执行
            for i, tool_name in enumerate(tool_chain.tools):
                print(f"执行工具链 {chain_id} 中的工具 {i+1}/{len(tool_chain.tools)}: {tool_name}")
                
                # 准备输入数据
                if i == 0:
                    tool_input = current_data
                else:
                    # 从前一个工具的输出获取数据
                    prev_tool = tool_chain.tools[i-1]
                    if prev_tool in tool_chain.data_flow and tool_chain.data_flow[prev_tool] == tool_name:
                        tool_input = self._transform_data(results[prev_tool]['result'], prev_tool, tool_name)
                    else:
                        tool_input = current_data
                
                # 执行工具
                result = await self.execute_tool(tool_name, tool_input, context)
                results[tool_name] = result
                
                if not result['success']:
                    raise Exception(f"工具 {tool_name} 执行失败: {result['error']}")
                
                current_data = result['result']
            
            # 更新工具链性能
            execution_time = (datetime.now() - start_time).total_seconds()
            tool_chain.success_rate = tool_chain.success_rate * 0.9 + 0.1
            tool_chain.average_execution_time = (
                tool_chain.average_execution_time * 0.9 + execution_time * 0.1
            )
            
            return {
                'success': True,
                'chain_id': chain_id,
                'results': results,
                'final_result': current_data,
                'execution_time': execution_time
            }
            
        except Exception as e:
            execution_time = (datetime.now() - start_time).total_seconds()
            tool_chain.success_rate = tool_chain.success_rate * 0.9
            
            return {
                'success': False,
                'chain_id': chain_id,
                'error': str(e),
                'execution_time': execution_time,
                'partial_results': results if 'results' in locals() else {}
            }
    
    def _transform_data(self, data: Any, source_tool: str, target_tool: str) -> Any:
        """转换数据格式"""
        # 简化的数据转换逻辑
        if source_tool == "enhanced_literature_analyzer" and target_tool == "smart_experiment_designer":
            # 文献分析结果转换为实验设计输入
            if isinstance(data, dict) and 'query' in data:
                return {
                    'context': self._string_to_research_context(data['query']),
                    'literature_analysis': data
                }
        
        elif source_tool == "smart_experiment_designer" and target_tool == "hf_data_search":
            # 实验设计结果转换为数据集搜索查询
            if isinstance(data, dict) and 'research_question' in data:
                return data['research_question']
        
        # 默认返回原数据
        return data
    
    def _string_to_research_context(self, query: str) -> ResearchContext:
        """字符串转换为研究上下文"""
        return ResearchContext(
            topic=query,
            phase="literature review",
            objectives=[f"研究{query}相关问题"],
            constraints={},
            resources={},
            timeline={}
        )
    
    def _paper_list_to_literature_analysis(self, papers: List[str]) -> Dict[str, Any]:
        """论文列表转换为文献分析"""
        return {
            'papers': papers,
            'total_papers': len(papers),
            'analysis_timestamp': datetime.now().isoformat()
        }
    
    def _dataset_list_to_data_requirements(self, datasets: List[Dict[str, Any]]) -> Dict[str, Any]:
        """数据集列表转换为数据需求"""
        return {
            'available_datasets': datasets,
            'recommended_dataset': datasets[0] if datasets else None,
            'data_requirements': {
                'size': 'medium',
                'type': 'structured',
                'quality': 'high'
            }
        }
    
    def get_tool_recommendations(self, task_description: str) -> List[Dict[str, Any]]:
        """获取工具推荐"""
        recommendations = []
        
        # 基于任务描述推荐工具
        task_lower = task_description.lower()
        
        if 'literature' in task_lower or 'paper' in task_lower or 'review' in task_lower:
            recommendations.append({
                'tool_name': 'enhanced_literature_analyzer',
                'confidence': 0.9,
                'reason': '任务涉及文献分析'
            })
            recommendations.append({
                'tool_name': 'arxiv_search',
                'confidence': 0.8,
                'reason': '需要搜索学术论文'
            })
        
        if 'experiment' in task_lower or 'design' in task_lower:
            recommendations.append({
                'tool_name': 'smart_experiment_designer',
                'confidence': 0.9,
                'reason': '任务涉及实验设计'
            })
        
        if 'data' in task_lower or 'dataset' in task_lower:
            recommendations.append({
                'tool_name': 'hf_data_search',
                'confidence': 0.8,
                'reason': '任务涉及数据集搜索'
            })
        
        if 'code' in task_lower or 'program' in task_lower or 'execute' in task_lower:
            recommendations.append({
                'tool_name': 'code_executor',
                'confidence': 0.7,
                'reason': '任务涉及代码执行'
            })
        
        # 按置信度排序
        recommendations.sort(key=lambda x: x['confidence'], reverse=True)
        
        return recommendations
    
    def get_tool_chain_recommendations(self, task_description: str) -> List[Dict[str, Any]]:
        """获取工具链推荐"""
        recommendations = []
        
        task_lower = task_description.lower()
        
        if 'literature review' in task_lower or ('literature' in task_lower and 'analysis' in task_lower):
            recommendations.append({
                'chain_id': 'literature_review_chain',
                'confidence': 0.9,
                'reason': '任务是文献综述'
            })
        
        if 'experiment' in task_lower and 'design' in task_lower:
            recommendations.append({
                'chain_id': 'experiment_design_chain',
                'confidence': 0.9,
                'reason': '任务是实验设计'
            })
        
        if 'data analysis' in task_lower or ('data' in task_lower and 'analysis' in task_lower):
            recommendations.append({
                'chain_id': 'data_analysis_chain',
                'confidence': 0.8,
                'reason': '任务是数据分析'
            })
        
        return recommendations
    
    def get_performance_report(self) -> Dict[str, Any]:
        """获取性能报告"""
        report = {
            'summary': dict(self.performance_stats),
            'tool_details': {},
            'tool_chain_performance': {},
            'recommendations': []
        }
        
        # 工具详细信息
        for tool_name, metadata in self.tool_metadata.items():
            report['tool_details'][tool_name] = {
                'usage_count': metadata.usage_count,
                'last_used': metadata.last_used.isoformat() if metadata.last_used else None,
                'performance_metrics': metadata.performance_metrics,
                'error_rate': self.performance_stats['error_rates'].get(tool_name, 0.0)
            }
        
        # 工具链性能
        for chain_id, chain in self.tool_chains.items():
            report['tool_chain_performance'][chain_id] = {
                'success_rate': chain.success_rate,
                'average_execution_time': chain.average_execution_time,
                'tools_count': len(chain.tools)
            }
        
        # 生成优化建议
        report['recommendations'] = self._generate_optimization_recommendations()
        
        return report
    
    def _generate_optimization_recommendations(self) -> List[str]:
        """生成优化建议"""
        recommendations = []
        
        # 基于错误率的建议
        for tool_name, error_rate in self.performance_stats['error_rates'].items():
            if error_rate > 0.1:  # 错误率超过10%
                recommendations.append(f"工具 {tool_name} 错误率较高 ({error_rate:.1%})，建议检查输入数据格式")
        
        # 基于使用频率的建议
        usage_freq = self.performance_stats['tool_usage_frequency']
        if usage_freq:
            most_used = max(usage_freq, key=usage_freq.get)
            least_used = min(usage_freq, key=usage_freq.get)
            
            if usage_freq[most_used] > usage_freq[least_used] * 10:
                recommendations.append(f"工具使用不均衡，考虑推广 {least_used} 的使用")
        
        # 基于执行时间的建议
        for tool_name, metadata in self.tool_metadata.items():
            avg_time = metadata.performance_metrics.get('average_execution_time', 0)
            if avg_time > 30:  # 超过30秒
                recommendations.append(f"工具 {tool_name} 执行时间较长，考虑优化或缓存")
        
        return recommendations
    
    def export_tool_configuration(self) -> Dict[str, Any]:
        """导出工具配置"""
        config = {
            'tools': {},
            'tool_chains': {},
            'data_transformers': list(self.data_transformers.keys()),
            'export_timestamp': datetime.now().isoformat()
        }
        
        # 导出工具信息
        for tool_name, metadata in self.tool_metadata.items():
            config['tools'][tool_name] = {
                'name': metadata.name,
                'version': metadata.version,
                'description': metadata.description,
                'input_types': metadata.input_types,
                'output_types': metadata.output_types,
                'capabilities': metadata.capabilities
            }
        
        # 导出工具链信息
        for chain_id, chain in self.tool_chains.items():
            config['tool_chains'][chain_id] = {
                'name': chain.name,
                'tools': chain.tools,
                'data_flow': chain.data_flow
            }
        
        return config
    
    def validate_tool_compatibility(self, tool_names: List[str]) -> Dict[str, Any]:
        """验证工具兼容性"""
        validation_result = {
            'compatible': True,
            'issues': [],
            'suggestions': []
        }
        
        # 检查工具是否存在
        for tool_name in tool_names:
            if tool_name not in self.registered_tools:
                validation_result['compatible'] = False
                validation_result['issues'].append(f"工具 {tool_name} 未注册")
        
        # 检查数据类型兼容性
        for i in range(len(tool_names) - 1):
            current_tool = tool_names[i]
            next_tool = tool_names[i + 1]
            
            if current_tool in self.tool_metadata and next_tool in self.tool_metadata:
                current_outputs = self.tool_metadata[current_tool].output_types
                next_inputs = self.tool_metadata[next_tool].input_types
                
                # 简化的兼容性检查
                if not any(output_type in next_inputs for output_type in current_outputs):
                    validation_result['issues'].append(
                        f"工具 {current_tool} 的输出类型与 {next_tool} 的输入类型不兼容"
                    )
                    validation_result['suggestions'].append(
                        f"需要在 {current_tool} 和 {next_tool} 之间添加数据转换器"
                    )
        
        return validation_result